## 用于划分训练集、验证集。没有测试集了
import os
import random
import shutil

# 定义源文件夹和目标文件夹路径
source_images_dir = (
    r"/home/shuai/Downloads/dataset0/images"  # 替换为你的images文件夹路径
)
source_labels_dir = (
    r"/home/shuai/Downloads/dataset0/labels"  # 替换为你的labels文件夹路径
)
output_base_dir = r"/home/shuai/Downloads/dataset"  # 替换为输出的根目录路径

# 创建输出文件夹
train_images_dir = os.path.join(output_base_dir, "train/images")
train_labels_dir = os.path.join(output_base_dir, "train/labels")
val_images_dir = os.path.join(output_base_dir, "val/images")
val_labels_dir = os.path.join(output_base_dir, "val/labels")

for dir_path in [train_images_dir, train_labels_dir, val_images_dir, val_labels_dir]:
    os.makedirs(dir_path, exist_ok=True)

# 获取所有图片文件名（假设图片和标签文件名一致，且图片为.jpg格式，标签为.txt格式）
all_images = [f for f in os.listdir(source_images_dir) if f.endswith(".jpg")]
all_labels = [
    f.replace(".jpg", ".txt") for f in all_images
]  # 确保标签文件名与图片文件名对应

# 随机打乱文件顺序
data = list(zip(all_images, all_labels))
random.shuffle(data)

# 按照9:1的比例进行划分
train_size = int(len(data) * 0.9)

train_data = data[:train_size]
val_data = data[train_size:]


# 定义复制函数
def copy_data(data, target_images_dir, target_labels_dir):
    for image_file, label_file in data:
        # 复制图片
        shutil.copy(
            os.path.join(source_images_dir, image_file),
            os.path.join(target_images_dir, image_file),
        )
        # 复制标签
        shutil.copy(
            os.path.join(source_labels_dir, label_file),
            os.path.join(target_labels_dir, label_file),
        )


# 复制数据到各个文件夹
copy_data(train_data, train_images_dir, train_labels_dir)
copy_data(val_data, val_images_dir, val_labels_dir)

print("数据集划分完成！")
